package com.tcm.smarthealth.security;

import com.tcm.smarthealth.config.JwtProperties;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.io.Decoders;
import io.jsonwebtoken.security.Keys;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Instant;
import java.time.OffsetDateTime;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.util.Map;

@Component
public class JwtTokenProvider {

    private static final Logger log = LoggerFactory.getLogger(JwtTokenProvider.class);
    private static final int MIN_KEY_BYTES = 32; // 256 bits

    private final JwtProperties jwtProperties;
    private Key signingKey;

    public JwtTokenProvider(JwtProperties jwtProperties) {
        this.jwtProperties = jwtProperties;
    }

    @PostConstruct
    public void init() {
        byte[] keyBytes = resolveKeyMaterial(jwtProperties.getSecretKey());
        this.signingKey = Keys.hmacShaKeyFor(keyBytes);
    }

    public String generateToken(Long userId, String username, String role) {
        Instant now = Instant.now();
        Instant expiry = now.plus(jwtProperties.getExpirationMinutes(), ChronoUnit.MINUTES);

        return Jwts.builder()
                .setSubject(String.valueOf(userId))
                .setIssuedAt(Date.from(now))
                .setExpiration(Date.from(expiry))
                .addClaims(Map.of(
                        "username", username,
                        "role", role
                ))
                .signWith(signingKey, SignatureAlgorithm.HS256)
                .compact();
    }

    public Jws<Claims> validateAndParse(String token) {
        return Jwts.parserBuilder()
                .setSigningKey(signingKey)
                .build()
                .parseClaimsJws(token);
    }

    public OffsetDateTime getExpiryFromNow() {
        return OffsetDateTime.now().plusMinutes(jwtProperties.getExpirationMinutes());
    }

    private byte[] resolveKeyMaterial(String rawSecret) {
        String trimmed = rawSecret != null ? rawSecret.trim() : "";
        if (trimmed.isEmpty()) {
            throw new IllegalArgumentException("JWT secretKey must not be empty");
        }

        byte[] keyBytes;
        try {
            keyBytes = Decoders.BASE64.decode(trimmed);
        } catch (IllegalArgumentException ex) {
            keyBytes = trimmed.getBytes(StandardCharsets.UTF_8);
        }

        if (keyBytes.length < MIN_KEY_BYTES) {
            keyBytes = strengthenKey(keyBytes);
            log.warn("JWT secret key is shorter than {} bytes. Strengthened the key using SHA-256 hash.", MIN_KEY_BYTES);
        }

        return keyBytes;
    }

    private byte[] strengthenKey(byte[] weakKeyMaterial) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            return digest.digest(weakKeyMaterial);
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException("SHA-256 MessageDigest not available", e);
        }
    }
}

